import rclpy
import rosbag2_py
from rosidl_runtime_py.utilities import get_message
from rclpy.serialization import deserialize_message
from rclpy.time import Time
from semantic_nav.log_utils import get_logger
logger = get_logger()


class BagUtilsBase:
    def __init__(self, bag_file):
        storage_options = rosbag2_py.StorageOptions(uri=bag_file, storage_id='sqlite3')
        converter_options = rosbag2_py.ConverterOptions(input_serialization_format='cdr', output_serialization_format='cdr')
        self.reader = rosbag2_py.SequentialReader()
        self.reader.open(storage_options, converter_options)
        self.topic_types = {topic.name: topic.type for topic in self.reader.get_all_topics_and_types()}

    def has_topic(self, topic_name) -> bool:
        """
        Check if topic exists
        """
        return topic_name in self.topic_types

    def get_msg_type(self, topic_name):
        """
        Get message type
        """
        msg_type = get_message(self.topic_types[topic_name])
        return msg_type

    def process(self):
        while self.reader.has_next():
            (topic, data, timestamp) = self.reader.read_next()
            msg = deserialize_message(data, self.get_msg_type(topic))
            if not self.callback(topic, msg, timestamp):
                break

    def callback(self, topic, msg, timestamp):
        """
        Placeholder function that will be overridden in derived classes
        """
        return True


class BagTimeUtils(BagUtilsBase):
    def __init__(self, bag_file):
        super().__init__(bag_file)
        self.topic_name = None
        self.first_timestamp = None
        self.last_timestamp = None
        self.message_count = 0

    def get_first_timestamp_callback(self, topic, msg, timestamp):
        if topic == self.topic_name:
            self.first_timestamp = Time.from_msg(msg.header.stamp)
            return False
        return True

    def frequency_callback(self, topic, msg, timestamp):
        if topic == self.topic_name:
            time = Time.from_msg(msg.header.stamp)
            if self.first_timestamp is None:
                self.first_timestamp = time
                return True
            self.last_timestamp = time
            self.message_count += 1
        return True

    def get_first_timestamp(self, topic_name):
        self.topic_name = topic_name
        self.callback = self.get_first_timestamp_callback
        self.process()
        return self.first_timestamp

    def get_frequency(self, topic_name):
        self.topic_name = topic_name
        self.callback = self.frequency_callback
        self.process()
        if self.message_count == 0:
            logger.error(f"No messages found for topic '{topic_name}'.")
            return 0.0
        # convert from nano sceonds to seconds
        duration = (self.last_timestamp - self.first_timestamp).nanoseconds / 1e9
        frequency = self.message_count / duration if duration > 0.0 else 0.0
        return frequency
